"""
Case Type   : DATAVEC模块
Case Name   : sparsevec字段创建hnsw l1距离索引，索引选项ef_construction参数覆盖，参数不同，索引查询准确度验证
Create At   : 2050/4/14
Owner       : @peilinqian
Description :
    step1:创建3张表;
    step2:创建自定义插入函数;
    step3:表插入相同数据;
    step4:创建索引，指定ef_construction为16;
    step5:创建索引，指定ef_construction为25;
    step6:创建索引，指定ef_construction为200;
    step7:不同参数索引查询召回率;
Expect      :
    step1:创建3张表; expect:成功
    step2:创建自定义插入函数; expect:成功
    step3:表插入相同数据; expect:成功
    step4:创建索引，指定ef_construction为16; expect:成功
    step5:创建索引，指定ef_construction为25; expect:成功
    step6:创建索引，指定ef_construction为200; expect:成功
    step7:不同参数索引查询召回率; expect:ef_construction值越大，召回率越高
History     :
"""
import os
import random
import unittest
from decimal import Decimal

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Constant import Constant
from testcase.utils.Logger import Logger
from yat.test import Node


class DataVecTestCase(unittest.TestCase):
    def setUp(self):
        self.log = Logger()
        self.log.info(f'----{os.path.basename(__file__)}:start----')
        self.sh = CommonSH('PrimaryDbUser')
        self.pri_dbuser = Node(node='PrimaryDbUser')
        self.sta_1_sh = CommonSH('Standby1DbUser')
        self.constant = Constant()
        self.com = Common()
        self.comsql = CommonSQL('PrimaryDbUser')
        self.tb_name = 't_datavec_index_case0046'
        self.fun_insert = 'f_datavec_index_case0046'

    def test_main(self):
        step_txt = '----step1:创建3张表; expect:成功----'
        self.log.info(step_txt)
        create_tb = f"drop table if exists {self.tb_name}_1,{self.tb_name}_2," \
                    f"{self.tb_name}_3;" \
                    f"create table {self.tb_name}_1 " \
                    f"(id bigserial primary key, val sparsevec(50));" \
                    f"create table {self.tb_name}_2 " \
                    f"(id bigserial primary key, val sparsevec(50));" \
                    f"create table {self.tb_name}_3 " \
                    f"(id bigserial primary key, val sparsevec(50));"
        result = self.sh.execut_db_sql(create_tb)
        self.log.info(result)
        self.assertIn(self.constant.TABLE_CREATE_SUCCESS, result,
                      '执行失败:' + step_txt)

        step_txt = '----step2:创建自定义插入函数; expect:成功----'
        self.log.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        self.log.info(result)
        self.assertIn(self.fun_insert, result, '执行失败:' + step_txt)

        step_txt = '----step3:表插入相同数据; expect:成功----'
        self.log.info(step_txt)
        insert_sql = f"select {self.fun_insert}(50,5000,'{self.tb_name}_1'," \
                     f"'val','sparsevec');"
        result = self.sh.execut_db_sql(insert_sql)
        self.log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        result = self.sh.execut_db_sql(insert_sql)
        self.log.info(result)
        self.assertIn('SUCCESS', result, '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_2 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        self.log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)
        sql = f"insert into {self.tb_name}_3 select * from {self.tb_name}_1;"
        result = self.sh.execut_db_sql(sql)
        self.log.info(result)
        self.assertIn(self.constant.INSERT_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step4:创建索引，指定ef_construction为16; expect:成功----'
        self.log.info(step_txt)
        create_index = f"set maintenance_work_mem='100MB';" \
                       f"create index idx_{self.tb_name}_1 on " \
                       f"{self.tb_name}_1 using hnsw(val sparsevec_l1_ops) " \
                       f"with(ef_construction=16,m=8);"
        result = self.sh.execut_db_sql(create_index)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step5:创建索引，指定ef_construction为25; expect:成功----'
        self.log.info(step_txt)
        create_index = f"set maintenance_work_mem='100MB';" \
                       f"create index idx_{self.tb_name}_2 on " \
                       f"{self.tb_name}_2 using hnsw(val sparsevec_l1_ops) " \
                       f"with(ef_construction=25,m=8)"
        result = self.sh.execut_db_sql(create_index)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step6:创建索引，指定ef_construction为200，大于默认值64; expect:成功----'
        self.log.info(step_txt)
        create_index = f"set maintenance_work_mem='100MB';" \
                       f"create index idx_{self.tb_name}_3 on " \
                       f"{self.tb_name}_3 using hnsw(val sparsevec_l1_ops) " \
                       f"with(ef_construction=200,m=8);"
        result = self.sh.execut_db_sql(create_index)
        self.log.info(result)
        self.assertIn(self.constant.CREATE_INDEX_SUCCESS_MSG, result,
                      '执行失败:' + step_txt)

        step_txt = '----step7:不同参数索引查询召回率; expect:ef_construction值越大，召回率越高----'
        self.log.info(step_txt)
        result1 = []
        result2 = []
        result3 = []
        for i in range(1, 100):
            self.log.info(f'----第{i}次查询----')
            radom_vector = random.randint(1, 10000)
            sql = f"select id, val<+>(select val from {self.tb_name}_1 where" \
                  f" id ={radom_vector}) from {self.tb_name}_1 " \
                  f"order by val <+> (select val from {self.tb_name}_1 where" \
                  f" id ={radom_vector}) limit 10;"
            result1.append(self.comsql.vector_query_recall(
                sql, 'set hnsw_ef_search=5;', is_detail=False))
            sql = f"select id, val <+>(select val from {self.tb_name}_2 where" \
                  f" id ={radom_vector}) from {self.tb_name}_2 " \
                  f"order by val <+> (select val from {self.tb_name}_2 where" \
                  f" id ={radom_vector}) limit 10;"
            result2.append(self.comsql.vector_query_recall(
                sql, 'set hnsw_ef_search=5;', is_detail=False))
            sql = f"select id, val <+>(select val from {self.tb_name}_3 where" \
                  f" id ={radom_vector}) from {self.tb_name}_3 " \
                  f"order by val<+>(select val from {self.tb_name}_3 where" \
                  f" id ={radom_vector}) limit 10;"
            result3.append(self.comsql.vector_query_recall(
                sql, 'set hnsw_ef_search=5;', is_detail=False))
        self.log.info(result1)
        total = sum(Decimal(str(x)) for x in result1)
        average1 = total / len(result1)
        self.log.info(average1)

        self.log.info(result2)
        total = sum(Decimal(str(x)) for x in result2)
        average2 = total / len(result2)
        self.log.info(average2)

        self.log.info(result3)
        total = sum(Decimal(str(x)) for x in result3)
        average3 = total / len(result3)
        self.log.info(average3)
        self.assertGreaterEqual(average2, average1, '执行失败:' + step_txt)
        self.assertGreaterEqual(average3, average2, '执行失败:' + step_txt)

    def tearDown(self):
        step_txt = '----this is teardown----'
        self.log.info(step_txt)
        step1_txt = '----teardown:清理测试数据; expect:成功----'
        self.log.info(step1_txt)
        drop_cmd = f"drop table if exists {self.tb_name}_1,{self.tb_name}_2," \
                   f"{self.tb_name}_3;"
        drop_result1 = self.sh.execut_db_sql(drop_cmd)
        self.log.info(drop_result1)
        drop_cmd = f"drop function if exists {self.fun_insert};"
        drop_result2 = self.sh.execut_db_sql(drop_cmd)
        self.log.info(drop_result2)

        self.log.info(f'----{os.path.basename(__file__)}:end----')
        self.assertIn(self.constant.TABLE_DROP_SUCCESS, drop_result1,
                      '执行失败:' + step_txt + step1_txt)
        self.assertIn(self.constant.DROP_FUNCTION_SUCCESS_MSG, drop_result2,
                      '执行失败:' + step_txt + step1_txt)
